# -*- coding: utf-8 -*-
"""
Created on Sat Feb  23 10:39:20 2013

@author: swLi
"""

import numpy as np
from numpy import linalg as LA

from traits.api import Int, HasTraits, List, Array, Float
from phy.interaction.MultiSpinInteraction import MultiSpinInteraction
from phy.interaction.SingleSpinInteraction import SingleSpinInteraction
from phy.general.MaterialData import DIAMOND_LATTICE_CONSTANT as DiamLattConst
from phy.general.MaterialData import DIAMOND_UNIT_CELL_POS as DiamUCPos
from statsmodels.sandbox.distributions.tests.check_moments import distlow

class Lattice(HasTraits):
    """
    doc for lattice.
    Make sure each basis are of form [x,y,z]
    """
    dim     = Int()
    basis   = List(Array)
    lattice_const      = Float()
    unit_cell_atom_num = Int()
    unit_cell_atom_pos = List(Array)
    
    def __init__(self, lattice_const = 1.0, basis = [ np.array([1.,0.,0.]), np.array([0.,1.,0.]), np.array([0.,0.,1.]) ],\
                  num = 1, pos = [np.array([0., 0., 0.])] ):
        self.basis         = basis
        self.dim           = len(self.basis)
        self.lattice_const = lattice_const
        self.unit_cell_atom_num = num
        self.unit_cell_atom_pos = pos
        
    def index2pos(self, idx, r):
        pos = (np.dot(idx, self.basis ) + self.unit_cell_atom_pos[r] ) * self.lattice_const
        return pos
    
    def pos2index(self, pos):
        return np.append(np.zeros(self.dim), 0)

class TooLargeRange(Exception):
    def __init__(self, value):
        self.value = value
    
class DiamondLattice(Lattice):
    impurity_num   = Int(0)
    impurity_index = List( Array )
    impurity_pos   = List( Array )
    
    def __init__(self, basis=[np.array([1, 0, 0]), np.array([0, 1, 0]), np.array([0, 0, 1])], num = 8, pos = DiamUCPos):
        Lattice.__init__(self, lattice_const = DiamLattConst, basis=basis, num = 8, pos = DiamUCPos)
        
    def impurity(self, abundance = 0.011, Rmax = 20, seednum = 1):
        """
        To generate a List of the positions of the impurities
        """
        num = int(Rmax ** self.dim * self.unit_cell_atom_num * abundance)
        
        np.random.seed(seednum)
        indices = [ tuple(np.random.randint(Rmax*2, size = self.dim) - Rmax) \
                    + ( np.random.randint(self.unit_cell_atom_num),) \
                   for i in range(num)]

        self.impurity_index = list(set(indices))
        self.impurity_num   = len(self.impurity_index)
        self.impurity_pos   = [ self.index2pos(index[:3], index[3]) for index in self.impurity_index]

    def impurity1(self, abundance = 0.011, Rmax = 10, seednum = 1):
        """
        a different realization
        """
        if Rmax > 20:
            raise TooLargeRange(Rmax)
        all_indices = [(i,j,k, r) for i in range(-Rmax, Rmax)\
                                  for j in range(-Rmax, Rmax)\
                                  for k in range(-Rmax, Rmax)\
                                  for r in range(self.unit_cell_atom_num)]
        
        np.random.seed(seednum)
        randlist = np.random.rand(len(all_indices))
        
        self.impurity_index = [idx for i, idx in enumerate(all_indices) if randlist[i] < abundance]#all_indices[randlist < abundance]                          
        self.impurity_num   = len(self.impurity_index)
        self.impurity_pos   = map(lambda idx: self.index2pos(idx[:3], idx[3]), self.impurity_index)
    
    def chop(self, dist= 10.0, relativeTo=np.array([0., 0., 0.])):
        distlist = [self.distance(pos) for pos in self.impurity_pos]
        self.impurity_pos   = [ self.impurity_pos[i] for i, d in enumerate(distlist) if d < dist]
        self.impurity_index = [ self.impurity_index[i] for i, d in enumerate(distlist) if d < dist]
        self.impurity_num   = len(self.impurity_index)
        
    def distance(self, pos, relativeTo=np.array([0., 0., 0.])):
        return LA.norm(pos-relativeTo)
        
    def add(self, pos=None, index=None):
        if index is not None:
            new_pos = self.index2pos(index[:3], index[3])
            new_idx = index
        elif pos is not None:
            new_pos = pos
            new_idx = self.pos2index(pos)
        else:
            print 'error'

        self.impurity_index.append(new_idx)
        self.impurity_pos.append(new_pos)
        self.impurity_num += 1

      
if __name__ == '__main__':
    lattice = DiamondLattice()
    lattice.impurity1(abundance=0.011,Rmax=10)
    lattice.chop(20.0)
    spins = lattice.impurity_pos
    print len(spins)
    print spins
    print 
    lattice.add(index=[1,1,1,0])
    print lattice.impurity_num
    print lattice.impurity_pos

